#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 增强版内存模式扫描器 - 专注于复杂Cookie模式识别

import re
import json
import os
from datetime import datetime

class PatternScanner:
    def __init__(self):
        # 常见cookie相关模式
        self.cookie_patterns = [
            # 基本的name=value格式
            r'([a-zA-Z0-9_\.\-]+)=([^;]+)',
            # 常见的Cookie头
            r'Cookie:\s*([^;\r\n]+)',
            # 特定平台TOKEN格式
            r'["\']?token["\']?\s*[:=]\s*["\']([^"\']+)["\']',
            # 应用特定的认证格式
            r'(access_token|auth_token|session_id)[=:]([^&"\';\s]+)',
            # JWT格式
            r'(eyJ[a-zA-Z0-9_-]+\.eyJ[a-zA-Z0-9_-]+\.[a-zA-Z0-9_-]+)',
            # 移动应用特有格式
            r'["\']X-Auth-Token["\']:\s*["\']([^"\']+)["\']',
            # 小程序格式
            r'(wx_login_token|mini_token)[=:]([^&"\';\s]+)',
        ]
        
        # 用户标识模式
        self.user_patterns = [
            r'uid[=:]([a-zA-Z0-9_\-]+)',
            r'userId[=:]([a-zA-Z0-9_\-]+)',
            r'user_id[=:]([a-zA-Z0-9_\-]+)',
            r'username[=:]([^@&"\';\s]+)',
            r'openid[=:]([a-zA-Z0-9_\-]+)',
            r'unionid[=:]([a-zA-Z0-9_\-]+)',
            r'email[=:]([^@\s]+@[^@\s]+\.[^@\s]+)',
            r'phone[=:](\d{11})',
        ]
        
        # 常见的Cookie域名
        self.common_domains = [
            'pinduoduo.com',
            'yangkeduo.com',
            'pdd.com',
            'mobile.yangkeduo.com',
            'api.pinduoduo.com',
            'api.yangkeduo.com',
        ]
        
        # 初始化结果存储
        self.results = {}
    
    def extract_user_id(self, text):
        """从文本中提取可能的用户ID"""
        for pattern in self.user_patterns:
            match = re.search(pattern, text, re.IGNORECASE)
            if match:
                return match.group(1)
        return "unknown"
    
    def scan_text(self, text):
        """扫描文本中的所有Cookie模式"""
        found_cookies = set()
        results = []
        
        # 提取可能的用户ID
        user_id = self.extract_user_id(text)
        
        # 扫描所有Cookie模式
        for pattern in self.cookie_patterns:
            matches = re.finditer(pattern, text, re.IGNORECASE)
            for match in matches:
                if len(match.groups()) >= 2:  # name=value格式
                    cookie = f"{match.group(1)}={match.group(2)}"
                else:  # 其他格式
                    cookie = match.group(0)
                
                # 检查是否已找到相同cookie
                if cookie not in found_cookies:
                    found_cookies.add(cookie)
                    
                    # 获取上下文
                    start_pos = max(0, match.start() - 50)
                    end_pos = min(len(text), match.end() + 50)
                    context = text[start_pos:end_pos]
                    
                    # 尝试识别域名
                    domain = "unknown"
                    for d in self.common_domains:
                        if d in text[max(0, match.start() - 100):match.end() + 100]:
                            domain = d
                            break
                    
                    results.append({
                        "cookie": cookie,
                        "user_id": user_id,
                        "domain": domain,
                        "context": context,
                        "pattern": pattern
                    })
        
        return results
    
    def classify_cookies(self, cookies):
        """对Cookie进行分类"""
        classified = {
            "authentication": [],
            "session": [],
            "tracking": [],
            "preference": [],
            "other": []
        }
        
        for cookie in cookies:
            cookie_str = cookie["cookie"].lower()
            
            # 认证相关
            if any(term in cookie_str for term in ['token', 'auth', 'login', 'access', 'jwt', 'id_']):
                classified["authentication"].append(cookie)
            
            # 会话相关
            elif any(term in cookie_str for term in ['session', 'sid', 'sess']):
                classified["session"].append(cookie)
            
            # 跟踪相关
            elif any(term in cookie_str for term in ['track', 'visitor', 'uid', 'analytics']):
                classified["tracking"].append(cookie)
            
            # 偏好相关
            elif any(term in cookie_str for term in ['pref', 'setting', 'config']):
                classified["preference"].append(cookie)
            
            # 其他
            else:
                classified["other"].append(cookie)
        
        return classified
    
    def add_results(self, user_id, cookies):
        """添加扫描结果到存储"""
        if user_id not in self.results:
            self.results[user_id] = []
        
        for cookie in cookies:
            if cookie["cookie"] not in [c["cookie"] for c in self.results[user_id]]:
                self.results[user_id].append(cookie)
    
    def save_results(self, output_dir="cookies_extracted"):
        """保存结果到文件"""
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 生成时间戳
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 保存完整结果
        full_path = os.path.join(output_dir, f"enhanced_cookies_{timestamp}.json")
        with open(full_path, 'w', encoding='utf-8') as f:
            json.dump(self.results, f, indent=2, ensure_ascii=False)
        
        print(f"[保存] 完整结果已保存到 {full_path}")
        
        # 保存分类结果
        for user_id, cookies in self.results.items():
            classified = self.classify_cookies(cookies)
            
            # 保存分类JSON
            class_path = os.path.join(output_dir, f"user_{user_id}_classified_{timestamp}.json")
            with open(class_path, 'w', encoding='utf-8') as f:
                json.dump(classified, f, indent=2, ensure_ascii=False)
            
            # 保存纯文本Cookie列表
            text_path = os.path.join(output_dir, f"user_{user_id}_cookies_{timestamp}.txt")
            with open(text_path, 'w', encoding='utf-8') as f:
                # 按重要性排序
                for category in ["authentication", "session", "tracking", "preference", "other"]:
                    if classified[category]:
                        f.write(f"=== {category.upper()} ===\n")
                        for cookie in classified[category]:
                            f.write(f"{cookie['cookie']}\n")
                        f.write("\n")
            
            print(f"[保存] 用户 {user_id} 的分类结果已保存")
        
        return full_path

# 示例用法
if __name__ == "__main__":
    scanner = PatternScanner()
    
    # 测试扫描
    test_text = """
    document.cookie = "session_id=abcdef123456; domain=example.com";
    xhr.setRequestHeader("Cookie", "token=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9; uid=12345");
    localStorage.setItem("auth_token", "secret_token_value");
    var userData = {userId: "user123", email: "test@example.com"};
    """
    
    results = scanner.scan_text(test_text)
    scanner.add_results("test_user", results)
    scanner.save_results() 